import gc
import json
import os
import random
import sys
import time
from functools import partial
from glob import glob
from pathlib import Path
import numpy as np

import deepspeed
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed as dist
from torch.utils.data import DataLoader
from torch.distributed.optim import ZeroRedundancyOptimizer
from torchvision.transforms import functional as TF
from diffusers import AutoencoderKL, DDPMScheduler, UNet2DConditionModel
from transformers import logging as tf_logging

from hydit.config_kolors import get_args
from hydit.lr_scheduler import WarmupLR
from hydit.data_loader.dataset_kolors_textmask import TextImageArrowStreamCustom
from hydit.diffusion import create_diffusion
from hydit.ds_config import deepspeed_config_from_args
from hydit.modules.ema import EMA
from hydit.modules.fp16_layers import Float16Module
from hydit.modules.posemb_layers import init_image_posemb
from hydit.utils.tools import create_logger, set_seeds, create_exp_folder, model_resume, get_trainable_params
from IndexKits.index_kits import ResolutionGroup
from IndexKits.index_kits.sampler import DistributedSamplerWithStartIndex, BlockDistributedSampler
from peft import LoraConfig, get_peft_model
from hydit.kolors_models.models.modeling_chatglm import ChatGLMModel
from hydit.kolors_models.models.tokenization_chatglm import ChatGLMTokenizer
from safetensors.torch import save_file

from peft.utils import get_peft_model_state_dict
from hydit.kolors_models.pipelines.pipeline_stable_diffusion_xl_chatglm_256 import StableDiffusionXLPipeline
from diffusers.utils import convert_state_dict_to_diffusers


def deepspeed_initialize(args, logger, model, opt, deepspeed_config):
    logger.info(f"Initialize deepspeed...")
    logger.info(f"    Using deepspeed optimizer")

    def get_learning_rate_scheduler(warmup_min_lr, lr, warmup_num_steps, opt):
        return WarmupLR(opt, warmup_min_lr, lr, warmup_num_steps)

    logger.info(f"    Building scheduler with warmup_min_lr={args.warmup_min_lr}, warmup_num_steps={args.warmup_num_steps}")
    model, opt, _, scheduler = deepspeed.initialize(
        model=model,
        model_parameters=get_trainable_params(model),
        config_params=deepspeed_config,
        args=args,
        lr_scheduler=partial(get_learning_rate_scheduler, args.warmup_min_lr, args.lr, args.warmup_num_steps) if args.warmup_num_steps > 0 else None,
    )
    return model, opt, scheduler

def save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoint_dir):

        # 处理参数名称的函数
    def modify_key_name(key):
        # 替换路径中的不需要的部分
        if 'base_model.model.' in key:
            key = key.replace('base_model.model.', '')
        return key

    def save_lora_weight(checkpoint_dir, client_state, tag=f"{train_steps:07d}.pt"):
        cur_ckpt_save_dir = f"{checkpoint_dir}/{tag}"
        if rank == 0:
            if args.use_fp16:
                # model.module.module.save_pretrained(cur_ckpt_save_dir)
                state_dict = get_peft_model_state_dict(model.module.module)
                modified_state_dict = {modify_key_name(key): value for key, value in state_dict.items()}
                unet_lora_layers_to_save = None
                unet_lora_layers_to_save = convert_state_dict_to_diffusers(modified_state_dict)
                StableDiffusionXLPipeline.save_lora_weights(cur_ckpt_save_dir, unet_lora_layers=unet_lora_layers_to_save)

            else:
                # model.module.save_pretrained(cur_ckpt_save_dir)
                unet_lora_layers_to_save = None
                unet_lora_layers_to_save = convert_state_dict_to_diffusers(get_peft_model_state_dict(model.module))
                StableDiffusionXLPipeline.save_lora_weights(cur_ckpt_save_dir,unet_lora_layers=unet_lora_layers_to_save)

    checkpoint_path = "[Not rank 0. Disabled output.]"

    client_state = {
        "steps": train_steps,
        "epoch": epoch,
        "args": args
    }
    if ema is not None:
        client_state['ema'] = ema.state_dict()

    dst_paths = []
    if train_steps % args.ckpt_every == 0:
        checkpoint_path = f"{checkpoint_dir}/{train_steps:07d}.pt"
        try:
            if args.training_parts == "lora":
                save_lora_weight(checkpoint_dir, client_state, tag=f"{train_steps:07d}.pt")
            else:
                model.save_checkpoint(checkpoint_dir, client_state=client_state, tag=f"{train_steps:07d}.pt")

            dst_paths.append(checkpoint_path)
            logger.info(f"Saved checkpoint to {checkpoint_path}")
        except:
            logger.error(f"Saved failed to {checkpoint_path}")

    if train_steps % args.ckpt_latest_every == 0 or train_steps == args.max_training_steps:
        save_name = "latest.pt"
        checkpoint_path = f"{checkpoint_dir}/{save_name}"
        try:
            if args.training_parts == "lora":
                save_lora_weight(checkpoint_dir, client_state, tag=f"{save_name}")
            else:
                model.save_checkpoint(checkpoint_dir, client_state=client_state, tag=f"{save_name}")
                
            dst_paths.append(checkpoint_path)
            logger.info(f"Saved checkpoint to {checkpoint_path}")
        except:
            logger.error(f"Saved failed to {checkpoint_path}")

    dist.barrier()
    if rank == 0 and len(dst_paths) > 0:
        # Delete optimizer states to avoid occupying too much disk space.
        for dst_path in dst_paths:
            for opt_state_path in glob(f"{dst_path}/zero_dp_rank_*_tp_rank_00_pp_rank_00_optim_states.pt"):
                os.remove(opt_state_path)

    return checkpoint_path


def prepare_model_inputs(args, batch_data, device):

    if args.is_textmask:
        image, textmask, text_embedding, text_embedding_mask, text_position_ids, image_meta_size, kwargs = batch_data
        textmask = textmask.to(device)
    else:
        image, text_embedding, text_embedding_mask, text_position_ids, image_meta_size, kwargs = batch_data
    # clip & mT5 text embedding

    image = image.to(device)
    text_embedding = text_embedding.to(device)
    text_embedding_mask = text_embedding_mask.to(device)
    text_position_ids = text_position_ids.to(device)
    # additional condition
    image_meta_size = image_meta_size.to(device)

    # Model conditions
    batch = {
        "pixel_values": image,
        "input_ids": text_embedding,
        "attention_mask":text_embedding_mask,
        "position_ids": text_position_ids,
        "add_time_ids":image_meta_size,
    }
    if args.is_textmask:
        batch['textmask'] = textmask
    return batch


def main(args):
    if args.training_parts == "lora":
        args.use_ema = False

    assert torch.cuda.is_available(), "Training currently requires at least one GPU."

    dist.init_process_group("nccl")
    world_size = dist.get_world_size()
    batch_size = args.batch_size
    grad_accu_steps = args.grad_accu_steps
    global_batch_size = world_size * batch_size * grad_accu_steps

    rank = dist.get_rank()
    device = rank % torch.cuda.device_count()
    seed = args.global_seed * world_size + rank
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.cuda.set_device(device)
    print(f"Starting rank={rank}, seed={seed}, world_size={world_size}.")
    deepspeed_config = deepspeed_config_from_args(args, global_batch_size)

    # Setup an experiment folder
    experiment_dir, checkpoint_dir, logger = create_exp_folder(args, rank)

    # Log all the arguments
    logger.info(sys.argv)
    logger.info(str(args))
    # Save to a json file
    args_dict = vars(args)
    args_dict['world_size'] = world_size
    with open(f"{experiment_dir}/args.json", 'w') as f:
        json.dump(args_dict, f, indent=4)

    # Disable the message "Some weights of the model checkpoint at ... were not used when initializing BertModel."
    # If needed, just comment the following line.
    tf_logging.set_verbosity_error()

    # ===========================================================================
    # Building Kolors
    # ===========================================================================

    logger.info("Building Kolors Model.")

    # ---------------------------------------------------------------------------
    #   Training sample base size, such as 256/512/1024. Notice that this size is
    #   just a base size, not necessary the actual size of training samples. Actual
    #   size of the training samples are correlated with `resolutions` when enabling
    #   multi-resolution training.
    # ---------------------------------------------------------------------------
    image_size = args.image_size
    if len(image_size) == 1:
        image_size = [image_size[0], image_size[0]]
    if len(image_size) != 2:
        raise ValueError(f"Invalid image size: {args.image_size}")
    assert image_size[0] % 8 == 0 and image_size[1] % 8 == 0, "Image size must be divisible by 8 (for the VAE encoder). " \
                                                              f"got {image_size}"
    latent_size = [image_size[0] // 8, image_size[1] // 8]

    # initialize model by deepspeed
    assert args.deepspeed, f"Must enable deepspeed in this script: train_deepspeed.py"
    with deepspeed.zero.Init(data_parallel_group=torch.distributed.group.WORLD,
                             remote_device=None if args.remote_device == 'none' else args.remote_device,
                             config_dict_or_path=deepspeed_config,
                             mpu=None,
                             enabled=args.zero_stage == 3):
        
        model = UNet2DConditionModel.from_pretrained(
            f'{args.pretrained_model_name_or_path}/unet',
            low_cpu_mem_usage=False,
            ignore_mismatched_sizes=True,
            revision=args.revision
        )

    if args.is_textmask:
        logger.info('Using textmask training!')
        

    # Multi-resolution / Single-resolution training.
    if args.multireso:
        resolutions = ResolutionGroup(image_size[0],
                                      align=16,
                                      step=args.reso_step,
                                      target_ratios=args.target_ratios).data
    else:
        resolutions = ResolutionGroup(image_size[0],
                                      align=16,
                                      target_ratios=['1:1']).data

    # freqs_cis_img = init_image_posemb(args.rope_img,
    #                                   resolutions=resolutions,
    #                                   patch_size=model.patch_size,
    #                                   hidden_size=model.hidden_size,
    #                                   num_heads=model.num_heads,
    #                                   log_fn=logger.info,
    #                                   rope_real=args.rope_real,
    #                                   )

    # Create EMA model and convert to fp16 if needed.
    ema = None
    if args.use_ema:
        ema = EMA(args, model, device, logger)

    if args.gradient_checkpointing:
        model.enable_gradient_checkpointing()

    # Setup FP16 main model:
    if args.use_fp16:
        model = Float16Module(model, args)
    logger.info(f"    Using main model with data type {'fp16' if args.use_fp16 else 'fp32'}")

    logger.info(f"    Loading vae")
    vae =  AutoencoderKL.from_pretrained(
        args.pretrained_model_name_or_path, subfolder="vae", revision=args.revision)
    
    logger.info(f"    Loading text encoder")
    text_encoder = ChatGLMModel.from_pretrained(
        args.pretrained_model_name_or_path,
        subfolder="text_encoder",
        revision=args.revision,
        trust_remote_code=True)

    logger.info(f"    Loading tokenizer")
    tokenizer = ChatGLMTokenizer.from_pretrained(
            os.path.join(args.pretrained_model_name_or_path, "text_encoder"),
            revision=args.revision,
            trust_remote_code=True)

    noise_scheduler = DDPMScheduler.from_pretrained(args.pretrained_model_name_or_path, subfolder="scheduler")
    
    if vae is not None:
        vae.requires_grad_(False)
    text_encoder.requires_grad_(False)

    if args.extra_fp16:
        weight_dtype = torch.float16
    else:
        weight_dtype = torch.float32

    model.to(device, dtype=weight_dtype)
    if vae is not None:
        vae.to(device, dtype=torch.float32)
    text_encoder.to(device, dtype=weight_dtype)


    logger.info(f"    Optimizer parameters: lr={args.lr}, weight_decay={args.weight_decay}")
    logger.info("    Using deepspeed optimizer")
    opt = None

    # ===========================================================================
    # Building Dataset
    # ===========================================================================

    logger.info(f"Building Streaming Dataset.")
    logger.info(f"    Loading index file {args.index_file} (v2)")

    dataset = TextImageArrowStreamCustom(args=args,
                                   resolution=image_size[0],
                                   random_flip=args.random_flip,
                                   log_fn=logger.info,
                                   index_file=args.index_file,
                                   multireso=args.multireso,
                                   batch_size=batch_size,
                                   world_size=world_size,
                                   random_shrink_size_cond=args.random_shrink_size_cond,
                                   merge_src_cond=args.merge_src_cond,
                                   uncond_p=args.uncond_p,
                                   text_ctx_len=args.text_len,
                                   tokenizer=tokenizer,
                                   )
    if args.multireso:
        sampler = BlockDistributedSampler(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed,
                                          shuffle=False, drop_last=True, batch_size=batch_size)
    else:
        sampler = DistributedSamplerWithStartIndex(dataset, num_replicas=world_size, rank=rank, seed=args.global_seed,
                                                   shuffle=False, drop_last=True)
    loader = DataLoader(dataset, batch_size=batch_size, sampler=sampler,
                        num_workers=args.num_workers, pin_memory=True, drop_last=True)
    logger.info(f"    Dataset contains {len(dataset):,} images.")
    logger.info(f"    Index file: {args.index_file}.")
    if args.multireso:
        logger.info(f'    Using MultiResolutionBucketIndexV2 with step {dataset.index_manager.step} '
                    f'and base size {dataset.index_manager.base_size}')
        logger.info(f'\n  {dataset.index_manager.resolutions}')

    # ===========================================================================
    # Loading parameter
    # ===========================================================================

    logger.info(f"Loading parameter")
    start_epoch = 0
    start_epoch_step = 0
    train_steps = 0
    # Resume checkpoint if needed
    
    # if args.resume is not None or len(args.resume) > 0:
    #     model, ema, start_epoch, start_epoch_step, train_steps = model_resume(args, model, ema, logger, len(loader))

    logger.info(f"    Training parts: {args.training_parts}")

    if args.training_parts == "lora":
        logger.info(f"Training part: lora")
        loraconfig = LoraConfig(
            r=args.rank,
            lora_alpha=args.rank,
            init_lora_weights="gaussian",
            target_modules=["to_k", "to_q", "to_v", "to_out.0"],
        )
        if args.use_fp16:
            model.module = get_peft_model(model.module, loraconfig)
            unet_lora_parameters = list(filter(lambda p: p.requires_grad, model.module.parameters()))
            total_params_bytes = sum(p.numel() * p.element_size() for p in unet_lora_parameters)
            total_params_mb = total_params_bytes / (1024 * 1024)
            logger.info(f'Total number of parameters with requires_grad=True: {total_params_mb:.2f} MB')
        else:
            model = get_peft_model(model, loraconfig)
            unet_lora_parameters = list(filter(lambda p: p.requires_grad, model.parameters()))
        
    
    model, opt, scheduler = deepspeed_initialize(args, logger, model, opt, deepspeed_config)

    # ===========================================================================
    # Training
    # ===========================================================================

    model.train()
    if args.use_ema:
        ema.eval()

    print(f"    Worker {rank} ready.")
    dist.barrier()

    iters_per_epoch = len(loader)
    logger.info(" ****************************** Running training ******************************")
    logger.info(f"      Number GPUs:               {world_size}")
    logger.info(f"      Number training samples:   {len(dataset):,}")
    logger.info(f"      Number parameters:         {sum(p.numel() for p in model.parameters()):,}")
    logger.info(f"      Number trainable params:   {sum(p.numel() for p in get_trainable_params(model)):,}")
    logger.info("    ------------------------------------------------------------------------------")
    logger.info(f"      Iters per epoch:           {iters_per_epoch:,}")
    logger.info(f"      Batch size per device:     {batch_size}")
    logger.info(f"      Batch size all device:     {batch_size * world_size * grad_accu_steps:,} (world_size * batch_size * grad_accu_steps)")
    logger.info(f"      Gradient Accu steps:       {args.grad_accu_steps}")
    logger.info(f"      Total optimization steps:  {args.epochs * iters_per_epoch // grad_accu_steps:,}")

    logger.info(f"      Training epochs:           {start_epoch}/{args.epochs}")
    logger.info(f"      Training epoch steps:      {start_epoch_step:,}/{iters_per_epoch:,}")
    logger.info(f"      Training total steps:      {train_steps:,}/{min(args.max_training_steps, args.epochs * iters_per_epoch):,}")
    logger.info("    ------------------------------------------------------------------------------")
    logger.info(f"      Noise schedule:            {args.noise_schedule}")
    logger.info(f"      Beta limits:               ({args.beta_start}, {args.beta_end})")
    logger.info(f"      Learn sigma:               {args.learn_sigma}")
    logger.info(f"      Prediction type:           {args.predict_type}")
    logger.info(f"      Noise offset:              {args.noise_offset}")

    logger.info("    ------------------------------------------------------------------------------")
    logger.info(f"      Using EMA model:           {args.use_ema} ({args.ema_dtype})")
    if args.use_ema:
        logger.info(f"      Using EMA decay:           {ema.max_value if args.use_ema else None}")
        logger.info(f"      Using EMA warmup power:    {ema.power if args.use_ema else None}")
    logger.info(f"      Using main model fp16:     {args.use_fp16}")
    logger.info(f"      Using extra modules fp16:  {args.extra_fp16}")
    logger.info("    ------------------------------------------------------------------------------")
    logger.info(f"      Experiment directory:      {experiment_dir}")
    logger.info("    *******************************************************************************")

    if args.gc_interval > 0:
        gc.disable()
        gc.collect()

    # Variables for monitoring/logging purposes:
    log_steps = 0
    running_loss = 0
    start_time = time.time()

    if args.async_ema:
        ema_stream = torch.cuda.Stream()

    # Training loop
    for epoch in range(start_epoch, args.epochs):
        logger.info(f"    Start random shuffle with seed={seed}")
        # Makesure all processors use the same seed to shuffle dataset.
        #dataset.shuffle(seed=args.global_seed + epoch, fast=True)
        logger.info(f"    End of random shuffle")

        # Move sampler to start_index
        if not args.multireso:
            start_index = start_epoch_step * world_size * batch_size
            if start_index != sampler.start_index:
                sampler.start_index = start_index
                # Reset start_epoch_step to zero, to ensure next epoch will start from the beginning.
                start_epoch_step = 0
                logger.info(f"      Iters left this epoch: {len(loader):,}")

        logger.info(f"    Beginning epoch {epoch}...")
        step = 0
        for batch_data in loader:
            step += 1

            batch = prepare_model_inputs(args, batch_data, device)
            # diffusion process
            pixel_values = batch["pixel_values"].to(dtype=weight_dtype)
            text_ids = batch['input_ids']
            attention_mask, position_ids = batch['attention_mask'], batch['position_ids']

            if vae is not None:
                # Convert images to latent space
                model_input = vae.encode(pixel_values.to(torch.float32)).latent_dist.sample()
                model_input = model_input.to(dtype=(weight_dtype))
                model_input = model_input * vae.config.scaling_factor

                if args.is_textmask:
                    _, _, h_latent, w_latent = model_input.shape
                    textmask = batch["textmask"]
                    textmask = torch.nn.functional.interpolate(textmask, size=(h_latent, w_latent)) # textmask下采样，倍数是VAE的下采样倍数
                    textmask = textmask.to(device)

            else:
                model_input = pixel_values

            # Sample noise that we'll add to the latents
            noise = torch.randn_like(model_input)
            noise += args.noise_offset * torch.randn(
                (model_input.shape[0], model_input.shape[1], 1, 1), device=model_input.device
            )

            bsz, channels, height, width = model_input.shape
            # Sample a random timestep for each image
            timesteps = torch.randint(
                0, noise_scheduler.config.num_train_timesteps, (bsz,), device=model_input.device
            )
            timesteps = timesteps.long()

            # Add noise to the model input according to the noise magnitude at each timestep
            # (this is the forward diffusion process)
            noisy_model_input = noise_scheduler.add_noise(model_input, noise, timesteps)

            output = text_encoder(
                input_ids=text_ids,
                attention_mask=attention_mask,
                position_ids=position_ids,
                output_hidden_states=True)
            prompt_embeds = output.hidden_states[-2].permute(1, 0, 2).clone() 
            text_proj = output.hidden_states[-1][-1, :, :].clone() 
        
            encoder_hidden_states = prompt_embeds

            add_time_ids_cond = batch['add_time_ids'].to(model_input.device)
            added_cond_kwargs = {"text_embeds": text_proj,
                                    "time_ids": add_time_ids_cond}

            # if accelerator.unwrap_model(super_model.unet).config.in_channels == channels * 2:
            #     noisy_model_input = torch.cat([noisy_model_input, noisy_model_input], dim=1)

            # Predict the noise residual
            model_pred = model(noisy_model_input, timesteps, encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs, return_dict=False)[0]
            
            if model_pred.shape[1] == 6:
                model_pred, _ = torch.chunk(model_pred, 2, dim=1)
            
            # Get the target for loss depending on the prediction type
            if noise_scheduler.config.prediction_type == "epsilon":
                target = noise
            elif noise_scheduler.config.prediction_type == "v_prediction":
                target = noise_scheduler.get_velocity(model_input, noise, timesteps)
            else:
                raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
            
            if args.is_textmask:
                # loss = F.mse_loss((model_pred*(1-textmask)).float(), (target*(1-textmask)).float(), reduction="mean")
                loss = F.mse_loss((model_pred*textmask).float(), (target*textmask).float(), reduction="mean")
    
            else:
                loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")


            # training model by deepspeed while use fp16
            if args.use_fp16:
                if args.use_ema and args.async_ema:
                    with torch.cuda.stream(ema_stream):
                        ema.update(model.module.module, step=step)
                    torch.cuda.current_stream().wait_stream(ema_stream)

            model.backward(loss)
            last_batch_iteration = (train_steps + 1) // (global_batch_size // (batch_size * world_size))
            model.step(lr_kwargs={'last_batch_iteration': last_batch_iteration})

            if args.use_ema and not args.async_ema or (args.async_ema and step == len(loader)-1):
                if args.use_fp16:
                    ema.update(model.module.module, step=step)
                else:
                    ema.update(model.module, step=step)

            # ===========================================================================
            # Log loss values:
            # ===========================================================================
            running_loss += loss.item()
            log_steps += 1
            train_steps += 1
            if train_steps % args.log_every == 0:
                # Measure training speed:
                torch.cuda.synchronize()
                end_time = time.time()
                steps_per_sec = log_steps / (end_time - start_time)
                # Reduce loss history over all processes:
                avg_loss = torch.tensor(running_loss / log_steps, device=device)
                dist.all_reduce(avg_loss, op=dist.ReduceOp.SUM)
                avg_loss = avg_loss.item() / world_size
                # get lr from deepspeed fused optimizer
                logger.info(f"(step={train_steps:07d}) " +
                            (f"(update_step={train_steps // args.grad_accu_steps:07d}) " if args.grad_accu_steps > 1 else "") +
                            f"Train Loss: {avg_loss:.4f}, "
                            f"Lr: {opt.param_groups[0]['lr']:.6g}, "
                            f"Steps/Sec: {steps_per_sec:.2f}, "
                            f"Samples/Sec: {int(steps_per_sec * batch_size * world_size):d}")
                # Reset monitoring variables:
                running_loss = 0
                log_steps = 0
                start_time = time.time()

            # collect gc:
            if args.gc_interval > 0 and (step % args.gc_interval == 0):
                gc.collect()

            if (train_steps % args.ckpt_every == 0 or train_steps % args.ckpt_latest_every == 0  # or train_steps == args.max_training_steps
                ) and train_steps > 0:
                save_checkpoint(args, rank, logger, model, ema, epoch, train_steps, checkpoint_dir)

            if train_steps >= args.max_training_steps:
                logger.info(f"Breaking step loop at {train_steps}.")
                break

        if train_steps >= args.max_training_steps:
            logger.info(f"Breaking epoch loop at {epoch}.")
            break

    dist.destroy_process_group()


if __name__ == "__main__":
    # Start
    main(get_args())
